#!/usr/bin/env python3
"""
离线模型设置脚本
用于在网络受限环境下配置Stable Diffusion模型
"""

import os
import sys
import requests
from pathlib import Path

def download_model_files():
    """下载必要的模型文件到本地"""
    print("🔄 开始下载Stable Diffusion模型文件...")
    
    # 创建本地模型目录
    model_dir = Path("models/stable-diffusion-v1-4")
    model_dir.mkdir(parents=True, exist_ok=True)
    
    # 模型文件URL列表（使用镜像源）
    model_files = {
        "model_index.json": "https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/model_index.json",
        "scheduler/scheduler_config.json": "https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/scheduler/scheduler_config.json",
        "text_encoder/config.json": "https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/text_encoder/config.json",
        "tokenizer/tokenizer_config.json": "https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/tokenizer/tokenizer_config.json",
        "unet/config.json": "https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/unet/config.json",
        "vae/config.json": "https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/vae/config.json",
    }
    
    try:
        for file_path, url in model_files.items():
            local_path = model_dir / file_path
            local_path.parent.mkdir(parents=True, exist_ok=True)
            
            print(f"📥 下载: {file_path}")
            response = requests.get(url, timeout=30)
            response.raise_for_status()
            
            with open(local_path, 'wb') as f:
                f.write(response.content)
            
            print(f"✅ 完成: {file_path}")
        
        print("\n🎉 配置文件下载完成！")
        print("⚠️  注意: 大型模型文件(如.safetensors)需要手动下载")
        print("📋 请访问以下链接手动下载:")
        print("   - https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/unet/diffusion_pytorch_model.safetensors")
        print("   - https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/vae/diffusion_pytorch_model.safetensors")
        print("   - https://hf-mirror.com/CompVis/stable-diffusion-v1-4/resolve/main/text_encoder/model.safetensors")
        
    except Exception as e:
        print(f"❌ 下载失败: {e}")
        return False
    
    return True

def setup_local_config():
    """设置本地模型配置"""
    config_content = '''
model:
  model_id: "./models/stable-diffusion-v1-4"  # 使用本地路径
  device: "cpu"
  torch_dtype: "float32"
  use_safetensors: true
  enable_attention_slicing: true
  enable_cpu_offload: false

generation:
  width: 512
  height: 512
  num_inference_steps: 20
  guidance_scale: 7.5
  num_images_per_prompt: 1
  seed: null

default_prompts:
  positive: "high quality, detailed, masterpiece"
  negative: "blurry, low quality, distorted, ugly"

output:
  directory: "outputs"
  format: "jpg"
  quality: 95
  filename_template: "{prompt}_{timestamp}"

performance:
  enable_memory_efficient_attention: true
  enable_sequential_cpu_offload: false
  enable_model_cpu_offload: false
  enable_attention_slicing: true
  slice_size: "auto"

logging:
  level: "INFO"
  format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
  file: "logs/auto_video_{date}.log"
'''
    
    with open("config/config_local.yaml", "w", encoding="utf-8") as f:
        f.write(config_content)
    
    print("✅ 本地配置文件已创建: config/config_local.yaml")

def main():
    print("🚀 Stable Diffusion 离线模型设置工具")
    print("=" * 50)
    
    choice = input("选择操作:\n1. 下载模型配置文件\n2. 创建本地配置\n3. 全部执行\n请输入选择 (1-3): ")
    
    if choice in ['1', '3']:
        if not download_model_files():
            print("❌ 模型文件下载失败")
            return
    
    if choice in ['2', '3']:
        setup_local_config()
    
    print("\n📝 使用说明:")
    print("1. 手动下载大型模型文件到 models/stable-diffusion-v1-4/ 目录")
    print("2. 使用命令: python main.py --config config/config_local.yaml generate \"your prompt\"")
    print("3. 或者修改默认配置文件指向本地模型路径")

if __name__ == "__main__":
    main()